import os
import paramiko
import stat
import time
import psutil

def check_disk_space(path, required_space):
    """
    检查指定路径下的磁盘空间是否足够
    :param path: 要检查的路径
    :param required_space: 需要的空间大小（以字节为单位）
    :return: True表示空间足够，False表示空间不足
    """
    disk_usage = psutil.disk_usage(path)
    available_space = disk_usage.free
    if available_space >= 10*1024*1024*1024:
        return True
    else:
        return False

def wait_for_disk_space(path, required_space, interval=10):
    """
    等待指定路径下的磁盘空间足够
    :param path: 要检查的路径
    :param required_space: 需要的空间大小（以字节为单位）
    :param interval: 检查间隔时间（秒）
    """
    while not check_disk_space(path, required_space):
        print("Waiting for disk space...")
        time.sleep(interval)

def scp(remote_host, remote_username, private_key, remote_dict, local_dict):
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        key = paramiko.RSAKey.from_private_key_file(private_key)
        client.connect(remote_host, username=remote_username, pkey=key)
        scp_client = client.open_sftp()
        scp_client.chdir(remote_dict)
        remote_files = scp_client.listdir()
        print(remote_files)
        for day in remote_files:
            if not stat.S_ISDIR(scp_client.lstat(remote_dict+ "/" + day).st_mode):
                continue
            scp_client.chdir(remote_dict+ "/" + day)
            clips = scp_client.listdir()
            for clip in clips:
                if not stat.S_ISDIR(scp_client.lstat(remote_dict+ "/" + day + "/"+ clip).st_mode):
                    continue
                scp_client.chdir(remote_dict+ "/" + day + "/"+ clip)
                files = scp_client.listdir()
                for file in files:
                    if file.endswith(".mkv"):
                        local_data_path = os.path.join(local_dict, day, clip)
                        if os.path.isdir(local_data_path):
                            if "vis" in os.listdir(local_data_path) or "depth" in os.listdir(local_data_path):
                                continue
                        if "updated" in files:
                            continue
                        remote_file_path = remote_dict+ "/" + day + "/"+ clip + "/"+ file
                        local_file_path = os.path.join(local_data_path, file)
                        print( remote_file_path, local_file_path)
                        required_space = scp_client.lstat(remote_file_path).st_size
                        if not check_disk_space(local_dict, required_space):
                            wait_for_disk_space(local_dict, required_space)
                        
                        os.makedirs(local_data_path, exist_ok=True)

                        start_time = time.time()
                        scp_client.get(remote_file_path, local_file_path)
                        end_time = time.time()
                        transferred_bytes = os.path.getsize(local_file_path)
                        transfer_time = end_time - start_time
                        transfer_speed = transferred_bytes / transfer_time
                        print("Transfer speed: {:.2f} MB/second".format(transfer_speed/1024/1024))
                        if "vis" in files:
                            with open(os.path.join(local_data_path, "vis"), "w"):
                                pass
                        else:
                            with open(os.path.join(local_data_path, "depth"), "w"):
                                pass
                        break
                    
                
    except paramiko.AuthenticationException:
        print("Authentication Fail")
    except paramiko.SSHException as e:
        print("SSH:", str(e))
    finally:
        if scp_client:
            scp_client.close()
        if client:
            client.close()

if __name__ == "__main__":
    HOST = "10.210.5.10"
    USER = "liuyun"
    PK_KEY = "C:\\Users\\liuyu\\.ssh\\id_rsa"
    REMOTE_DIR = "/share/datasets/hhodataset/azure5"
    LOCAL_DIR = "D:\\mkv\\azure5"
    scp(HOST, USER, PK_KEY, REMOTE_DIR, LOCAL_DIR)
        